#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np
import os

dtype_emu = {
    np.float16: 1,       # FLOAT16（半精度浮点数）
    np.float32: 2,       # FLOAT（单精度浮点数）
    np.float64: 3,       # DOUBLE（双精度浮点数）
    np.uint8: 4,        # UINT8（8位无符号整数）
    np.uint16: 5,        # UINT16（16位无符号整数）
    np.uint32: 6,        # UINT32（32位无符号整数）
    np.uint64: 7,        # UINT64（64位无符号整数）
    np.int8: 8,          # INT8（8位有符号整数）
    np.int16: 9,        # INT16（16位有符号整数）
    np.int32: 10,        # INT32（32位有符号整数）
    np.int64: 11         # INT64（64位有符号整数）
}

def custom_maskfill(input_x, mask, value):
    """自定义MaskFill"""
    output = np.copy(input_x)
    output[mask] = value
    return output


def gen_golden_data_maskfill():

    # 创建目录
    input_dir = "./input"
    output_dir = "./output"
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    dtype = np.float64

    ## 核间均分，单核计算量对齐:
    # input_shape = [32, 1024]

    ## 核间均分，单核计算量非对齐:
    # input_shape = [8, 1023]

    ## 核间不均分，单核计算量对齐:
    # input_shape = [32, 1023]

    ## 核间不均分，单核计算量非对齐:
    input_shape = [19, 37]
    
    # 数据乘积后向上对齐不得超过Addr最大存储196608

    mask_true_ratio = 0.4
    fill_value = -1.0
    
    # 生成数据
    input_x = np.random.uniform(-300, 300, input_shape).astype(dtype)
    mask = np.random.choice([False, True], size=input_shape, p=[1 - mask_true_ratio, mask_true_ratio]).astype(np.bool_)
    fill_value_casted = np.array(fill_value, dtype=dtype)

    total_elements = input_shape[0] * input_shape[1]
    tiling = np.array([total_elements, dtype_emu[dtype]], dtype=np.uint32)

    golden = custom_maskfill(input_x, mask, fill_value_casted)
    
    # 保存数据
    tiling.tofile(os.path.join(input_dir, "input_tiling.bin"))
    input_x.tofile(os.path.join(input_dir, "input_x.bin"))
    mask.tofile(os.path.join(input_dir, "mask.bin"))
    fill_value_casted.tofile(os.path.join(input_dir, "fill_value.bin"))
    golden.tofile(os.path.join(output_dir, "golden.bin"))

if __name__ == "__main__":
    gen_golden_data_maskfill()